% 1. RGB Tracing & Artery-Vein Separation Code
clear; clc;

% STEP 1: Load segmented image and corresponding RGB ground truth image
[segFile, segPath] = uigetfile('*', 'Select segmented vessel image');
if isequal(segFile, 0)
    disp('Selection cancelled for segmented image');
    return;
else
    segmentedImg = imread(fullfile(segPath, segFile));
    figure, imshow(segmentedImg), title('Segmented Vessel Image');
end

[rgbFile, rgbPath] = uigetfile('*', 'Select RGB ground truth image');
if isequal(rgbFile, 0)
    disp('Selection cancelled for RGB ground truth image');
    return;
else
    rgbImg = imread(fullfile(rgbPath, rgbFile));
    figure, imshow(rgbImg), title('RGB Ground Truth Image');
end

% STEP 2: Resize RGB image to match segmented image if necessary
[segRows, segCols, ~] = size(segmentedImg);
[rgbRows, rgbCols, ~] = size(rgbImg);

if segRows ~= rgbRows || segCols ~= rgbCols
    rgbImg = imresize(rgbImg, [segRows segCols], 'nearest');
    disp('Resized RGB ground truth image to match segmented image dimensions');
    figure, imshow(rgbImg), title('Resized RGB Image');
end

% STEP 3: Convert RGB image to HSV for color-based classification
hsvImg = rgb2hsv(rgbImg);
H = hsvImg(:,:,1);
S = hsvImg(:,:,2);
V = hsvImg(:,:,3);

% STEP 4: Define HSV thresholds for artery, vein, and overlap
arteryMask = ((H >= 0 & H <= 0.1) | (H >= 0.9 & H <= 1)) & (S > 0.4) & (V > 0.3);
veinMask = (H >= 0.55 & H <= 0.75) & (S > 0.4) & (V > 0.3);
overlapMask = (H >= 0.2 & H <= 0.45) & (S > 0.3) & (V > 0.3);

% STEP 5: Create binary masks for artery and vein based on segmentation and color labels
arteryImg = segmentedImg > 0 & arteryMask;
veinImg = segmentedImg > 0 & veinMask;

% Add overlap regions to both artery and vein to preserve connectivity
arteryImg = arteryImg | (overlapMask & (segmentedImg > 0));
veinImg = veinImg | (overlapMask & (segmentedImg > 0));

% STEP 6: Morphological closing to enhance connectivity
se = strel('disk', 2);
arteryImg = imclose(arteryImg, se);
veinImg = imclose(veinImg, se);

% STEP 7: Display and save final artery and vein masks
figure, imshow(uint8(arteryImg) * 255), title('Final Artery Mask');
imwrite(uint8(arteryImg) * 255, 'output_segmented_artery.png');

figure, imshow(uint8(veinImg) * 255), title('Final Vein Mask');
imwrite(uint8(veinImg) * 255, 'output_segmented_vein.png');
